#include "math.h"
#include <stdlib.h>
#include "inc/core/particles.h"
#include "inc/core/lib.h"
#include "stdio.h"

int compare(void const* one, void const* two)
{
    if (*(floatdig*)one > *(floatdig*)two)
        return 1;
    if (*(floatdig*)one < *(floatdig*)two)
        return -1;
    return 0;
}

void particle_initialize_particle_set(particles_set* set, int N, floatdig charge, floatdig mass)
{
    if (set != NULL)
    {
        if (N > 0)
        {
            particle *particles;
            set->particles = (particle*) malloc(sizeof (particle) * N);
            particles = set->particles;
            if (particles != NULL)
            {
                int i = 0;
                particle* one_particle = NULL;
                for (i = 0; i < N; i++)
                {
                    one_particle = &(particles[i]);
                    lib_set_vector(&(one_particle->pos), 0);
                    lib_set_vector(&(one_particle->speed), 0);
                    one_particle->active = PARTICLE_INACTIVE;
                }
            }
        }
        else
        {
            set->particles = NULL;
        }
        set->num_of_particles = N;
        set->charge = charge;
        set->mass = mass;
        gsl_rng_env_setup();
        set->r = gsl_rng_alloc(gsl_rng_taus);
    }
}

void particle_set_particle_speed_maxwellian(particles_set* set, int particle, int axe, floatdig sigma)
{
   set->particles[particle].speed.v[axe] = (floatdig) gsl_ran_gaussian(set->r, sigma);
}

void particle_clear_particle_set(particles_set* set)
{
    if (set != NULL)
    {
        gsl_rng_free(set->r);
        free(set->particles);
        set->particles = NULL;
        set->num_of_particles = 0;
    }
}

void particle_distribute_particle_set_speed_maxwell(particles_set* set, double sigma, int axe)
{
    if ((set != NULL) && (set->particles != NULL))
    {
        int i = 0;

        for (i = 0; i < set->num_of_particles; i++)
        {
            if (set->particles[i].active == PARTICLE_ACTIVE)
            {
                set->particles[i].speed.v[axe] += (floatdig) gsl_ran_gaussian(set->r, sigma);
            }
        }
    }
}

void particle_make_speed_disturbance(particles_set* set, floatdig as, floatdig k, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    speed3d* vec;
    for (i = 0; i < N; i++)
    {
        one_particle = &(set->particles[i]);
        vec = &(one_particle->speed);
        vec->v[axe] += as * cos(k * one_particle->pos.v[axe]);
    }
}

void particle_make_pos_disturbance(particles_set* set, floatdig ap, floatdig k, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    pos3d* pos;
    for (i = 0; i < N; i++)
    {
        one_particle = &(set->particles[i]);
        pos = &(one_particle->pos);
        pos->v[axe] += ap * sin(k * pos->v[axe]);
    }
}

void particle_distribute_particle_set_beam(particles_set* set, floatdig beam_speed, int axe)
{
    int i;
    int N = set->num_of_particles;
    particle* one_particle;
    speed3d* vec;
    for (i = 0; i < N; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            one_particle = &(set->particles[i]);
            vec = &(one_particle->speed);
            vec->v[axe] = beam_speed;
        }

    }
}

void particle_distrubute_particles_set_pos_uniform(particles_set* set, int axe, floatdig min_pos, floatdig max_pos)
{
    if ((set != NULL) && (set->particles != NULL))
    {
        int N = set->num_of_particles;

        floatdig delta = max_pos - min_pos;
        vector3d *vec = NULL;
        int i = 0;
        int j = 0;
        gsl_rng_env_setup();
        // select random number generator
        for (i = 0; i < set->num_of_particles; i++)
        {
            vec = &(set->particles[i].pos);
            vec->v[axe] = (floatdig) (delta * gsl_rng_uniform(set->r) + min_pos);
            set->particles[i].active = PARTICLE_ACTIVE;

        }
    }


}

particle* particle_get_particle(particles_set* set, int num)
{
    if ((set != NULL) && (set->particles != NULL) && (num < set->num_of_particles))
    {
        return &(set->particles[num]);
    }
    else
    {
        return NULL;
    }
}

void particle_put_particles_pos_into_file(char* fname, particles_set* set, int axe)
{
    FILE* to = fopen(fname, "w");
    if ((to != NULL) && (set != NULL) & (set->particles != NULL))
    {
        int i = 0;
        int N = set->num_of_particles;
        pos3d* pos = NULL;
        for (i = 0; i < N; i++)
        {
            if (set->particles[i].active == PARTICLE_ACTIVE)
            {
                pos = &(set->particles[i].pos);
                fprintf(to, "%d\t%f\n", i, (pos->v[axe]));
            }
        }
        fclose(to);
    }
}

void particle_put_particles_speed_into_file(char* fname, particles_set* set, int axe)
{
    FILE* to = fopen(fname, "w");
    if ((to != NULL) && (set != NULL) & (set->particles != NULL))
    {
        int i = 0;
        int N = set->num_of_particles;
        speed3d* speed = NULL;
        for (i = 0; i < N; i++)
        {
            speed = &(set->particles[i].speed);
            fprintf(to, "%d\t%e\n", i, (speed->v[axe]));
        }
        fclose(to);
    }
}

void particle_add_particle(particles_set* set, particle* one)
{
    if ((set != NULL) && (one != NULL))
    {
        if ((set->particles == NULL) || (set->num_of_particles == 0))
        {
            set->particles = (particle*) malloc(sizeof (particle));
            memcpy(&(set->particles[0]), one, sizeof (particle));
            set->num_of_particles = 1;
        }
        else
        {
            int i = 0;
            particle* new_set = (particle*) malloc(sizeof (particle)*(set->num_of_particles + 1));
            for (i = 0; i < set->num_of_particles; i++)
            {
                memcpy(&(new_set[i]), &(set->particles[i]), sizeof (particle));
            }
            memcpy(&(new_set[set->num_of_particles]), one, sizeof (particle));
            free(set->particles);
            set->particles = new_set;
            set->num_of_particles = set->num_of_particles + 1;
        }
    }
}

void particle_set_two_beams(particles_set* set, floatdig v01, floatdig sigma1,
                            floatdig v02, floatdig sigma2, floatdig min_z, floatdig max_z)
{
    gsl_rng * r;
    int i;
    int N = set->num_of_particles;
    int half = N / 2;
    floatdig sL = max_z - min_z;
    floatdig v_min = -5 * v01;
    floatdig v_max = 5 * v01;
    floatdig v;
    floatdig v_part = 0;

    gsl_rng_env_setup();
    // select random number generator
    r = gsl_rng_alloc(gsl_rng_ranlux389);
    v = v_min + (v_max - v_min) * gsl_rng_uniform(r);

    for (i = 0; i < half; i++)
    {
        v_part = set->particles[i].pos.v[Z] = sL * gsl_rng_uniform(r) + min_z;
        set->particles[i].speed.v[VZ] = v01 + gsl_ran_gaussian(r, sigma1);
        set->particles[i].active = PARTICLE_ACTIVE;

        set->particles[half + i].pos.v[Z] = sL * gsl_rng_uniform(r) + min_z;
        set->particles[half + i].speed.v[VZ] = v02 + gsl_ran_gaussian(r, sigma2);
        set->particles[half + i].active = PARTICLE_ACTIVE;
    }


}

void particle_dump_particles_pos_into_file(particles_set* set, char* path,
                                       char* filename, floatdig pos,
                                       floatdig dz)
{
    int i = 0;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    sprintf(fullname, "%s/%s", path, filename);
    to = fopen(fullname, "w");
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            if (dz > 0)
            {
                if (fabs(set->particles[i].pos.v[Z] - pos) < dz/2)
                {
                    fprintf(to, "%e\t%e\t%e\n", set->particles[i].pos.v[X], set->particles[i].pos.v[Y], set->particles[i].pos.v[Z]);
                }
            }
            else
            {
                fprintf(to, "%e\t%e\t%e\n", set->particles[i].pos.v[X], set->particles[i].pos.v[Y], set->particles[i].pos.v[Z]);
            }
        }
    }
    fclose(to);
}

void particle_dump_particles_speed_into_file(particles_set* set, char* path,
                                       char* filename, floatdig pos,
                                       floatdig dz)
{
    int i = 0;
    char fullname[MAX_BUF_SIZE];
    FILE* to;
    sprintf(fullname, "%s/%s", path, filename);
    to = fopen(fullname, "w");
    for (i = 0; i < set->num_of_particles; i++)
    {
        if (set->particles[i].active == PARTICLE_ACTIVE)
        {
            if (dz > 0)
            {
                if (fabs(set->particles[i].pos.v[Z] - pos) < dz/2)
                {
                    fprintf(to, "%e\t%e\t%e\t\n", set->particles[i].speed.v[VX], set->particles[i].speed.v[VY], set->particles[i].speed.v[VZ]);
                }
            }
            else
            {
                fprintf(to, "%e\t%e\t%e\t\n", set->particles[i].speed.v[VX], set->particles[i].speed.v[VY], set->particles[i].speed.v[VZ]);
            }
        }
    }
    fclose(to);
}


void particle_dump_trenergy_into_file(particles_set* set, char* path,
                                       char* filename, floatdig pos,
                                       floatdig dz, floatdig num_of_steps)
{
    int i = 0;
    int k = 0;
    char fullname[MAX_BUF_SIZE];
    int num_of_particles = 0;
    floatdig tre = 0;
    floatdig* vtr;
    floatdig min_vtr;
    floatdig max_vtr;
    floatdig vtr_step;
    floatdig v;
    FILE* to;
    sprintf(fullname, "%s/%s", path, filename);
    char* counted;
    to = fopen(fullname, "w");
    for (i = 0; i < set->num_of_particles; i++)
    {
        
        if ((set->particles[i].active == PARTICLE_ACTIVE) &&
            (fabs(set->particles[i].pos.v[Z] - pos) < dz/2))
        {
            //tre += 0.5*set->mass*(set->particles[i].speed.v[VX]*set->particles[i].speed.v[VX]+set->particles[i].speed.v[VY]*set->particles[i].speed.v[VY]);
            num_of_particles++;
        }
        
    }
    vtr = (floatdig*)malloc(sizeof(floatdig)*num_of_particles);
    counted = (char*)malloc(sizeof(char)*num_of_particles);
    for (i = 0; i < set->num_of_particles; i++)
    {

        if ((set->particles[i].active == PARTICLE_ACTIVE) &&
            (fabs(set->particles[i].pos.v[Z] - pos) < dz/2))
        {
            memcpy(&(vtr[k]), &(set->particles[i].trspeed), sizeof(floatdig));
            counted[k]=0;
            k++;
        }

    }
    qsort(vtr, num_of_particles, sizeof(floatdig) , compare);
    min_vtr = vtr[0];
    max_vtr = vtr[num_of_particles - 1];
    vtr_step = (max_vtr - min_vtr)/num_of_steps;
    for (v = min_vtr; v <= max_vtr; v += vtr_step)
    {
        int num_of_particles_with_the_speed = 0;
        for (k = 0; k < num_of_particles; k++)
        {
            if ((!counted[k])&&(fabs(v - vtr[k]) < vtr_step/2))
            {
                counted[k]=1;
                num_of_particles_with_the_speed++;
            }
        }
        fprintf(to, "%d\t%e\n", num_of_particles_with_the_speed, v);
        num_of_particles_with_the_speed = 0;
    }
    free(vtr);
    free(counted);
    fclose(to);
}
